/*
 * HIP DCU Correctness Validation Framework
 * Validates correctness of masquerading vs sequential execution
 * Tests performance improvement claims for 128 sub-tasks × 1024×1024 matrices
 */

#include <hip/hip_runtime.h>
#include <iostream>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iomanip>

#define CHECK_HIP(call) \
    do { \
        hipError_t error = call; \
        if (error != hipSuccess) { \
            std::cerr << "HIP error at " << __FILE__ << ":" << __LINE__ \
                      << " - " << hipGetErrorString(error) << std::endl; \
            exit(1); \
        } \
    } while(0)

// Simple vector addition kernel for validation
__global__ void vectorAddKernel(const float* a, const float* b, float* c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

// Masquerading kernel for multiple sub-tasks
__global__ void masqueradingKernel(float** a_arrays, float** b_arrays, float** c_arrays, 
                                  int arraySize, int numSubTasks) {
    int globalIdx = blockIdx.x * blockDim.x + threadIdx.x;
    int totalElements = arraySize * numSubTasks;
    
    if (globalIdx < totalElements) {
        int subTaskId = globalIdx / arraySize;
        int localIdx = globalIdx % arraySize;
        
        if (subTaskId < numSubTasks) {
            float* a = a_arrays[subTaskId];
            float* b = b_arrays[subTaskId];
            float* c = c_arrays[subTaskId];
            c[localIdx] = a[localIdx] + b[localIdx];
        }
    }
}

class CorrectnessValidator {
private:
    int arraySize = 1048576;  // 1024*1024 elements
    int numSubTasks = 128;
    
    // Host arrays
    float *h_a, *h_b, *h_c_sequential, *h_c_masquerading, *h_reference;
    
    // Device arrays
    float **d_a, **d_b, **d_c;  // Array of pointers for each sub-task
    float **d_a_ptrs, **d_b_ptrs, **d_c_ptrs;  // Device pointer arrays
    
public:
    CorrectnessValidator() {
        setupMemory();
    }
    
    ~CorrectnessValidator() {
        cleanup();
    }
    
    void setupMemory() {
        // Allocate host memory
        CHECK_HIP(hipHostMalloc(&h_a, arraySize * numSubTasks * sizeof(float)));
        CHECK_HIP(hipHostMalloc(&h_b, arraySize * numSubTasks * sizeof(float)));
        CHECK_HIP(hipHostMalloc(&h_c_sequential, arraySize * numSubTasks * sizeof(float)));
        CHECK_HIP(hipHostMalloc(&h_c_masquerading, arraySize * numSubTasks * sizeof(float)));
        CHECK_HIP(hipHostMalloc(&h_reference, arraySize * numSubTasks * sizeof(float)));
        
        // Initialize data
        initializeData();
        
        // Allocate device arrays for sub-tasks
        d_a = new float*[numSubTasks];
        d_b = new float*[numSubTasks];
        d_c = new float*[numSubTasks];
        
        for (int i = 0; i < numSubTasks; i++) {
            CHECK_HIP(hipMalloc(&d_a[i], arraySize * sizeof(float)));
            CHECK_HIP(hipMalloc(&d_b[i], arraySize * sizeof(float)));
            CHECK_HIP(hipMalloc(&d_c[i], arraySize * sizeof(float)));
        }
        
        // Setup device pointer arrays
        CHECK_HIP(hipMalloc(&d_a_ptrs, numSubTasks * sizeof(float*)));
        CHECK_HIP(hipMalloc(&d_b_ptrs, numSubTasks * sizeof(float*)));
        CHECK_HIP(hipMalloc(&d_c_ptrs, numSubTasks * sizeof(float*)));
        
        CHECK_HIP(hipMemcpy(d_a_ptrs, d_a, numSubTasks * sizeof(float*), hipMemcpyHostToDevice));
        CHECK_HIP(hipMemcpy(d_b_ptrs, d_b, numSubTasks * sizeof(float*), hipMemcpyHostToDevice));
        CHECK_HIP(hipMemcpy(d_c_ptrs, d_c, numSubTasks * sizeof(float*), hipMemcpyHostToDevice));
    }
    
    void initializeData() {
        for (int i = 0; i < arraySize * numSubTasks; i++) {
            h_a[i] = static_cast<float>(rand()) / RAND_MAX;
            h_b[i] = static_cast<float>(rand()) / RAND_MAX;
            h_reference[i] = h_a[i] + h_b[i];
        }
    }
    
    void runSequential() {
        const int gridSize = (arraySize + 255) / 256;
        const int threadsPerBlock = 256;
        
        std::cout << "Running sequential mode..." << std::endl;
        auto start = std::chrono::high_resolution_clock::now();
        
        // Copy all data to device
        for (int i = 0; i < numSubTasks; i++) {
            CHECK_HIP(hipMemcpy(d_a[i], &h_a[i * arraySize], arraySize * sizeof(float), hipMemcpyHostToDevice));
            CHECK_HIP(hipMemcpy(d_b[i], &h_b[i * arraySize], arraySize * sizeof(float), hipMemcpyHostToDevice));
        }
        
        // Execute sequentially
        for (int i = 0; i < numSubTasks; i++) {
            hipLaunchKernelGGL(vectorAddKernel, dim3(gridSize), dim3(threadsPerBlock), 
                             0, 0, d_a[i], d_b[i], d_c[i], arraySize);
            CHECK_HIP(hipDeviceSynchronize());
        }
        
        // Copy results back
        for (int i = 0; i < numSubTasks; i++) {
            CHECK_HIP(hipMemcpy(&h_c_sequential[i * arraySize], d_c[i], arraySize * sizeof(float), hipMemcpyDeviceToHost));
        }
        
        auto end = std::chrono::high_resolution_clock::now();
        double duration = std::chrono::duration<double, std::milli>(end - start).count();
        
        std::cout << "Sequential execution time: " << duration << " ms" << std::endl;
    }
    
    void runMasquerading() {
        const int totalElements = arraySize * numSubTasks;
        const int gridSize = (totalElements + 255) / 256;
        const int threadsPerBlock = 256;
        
        std::cout << "Running masquerading mode..." << std::endl;
        auto start = std::chrono::high_resolution_clock::now();
        
        // Copy all data to device
        for (int i = 0; i < numSubTasks; i++) {
            CHECK_HIP(hipMemcpy(d_a[i], &h_a[i * arraySize], arraySize * sizeof(float), hipMemcpyHostToDevice));
            CHECK_HIP(hipMemcpy(d_b[i], &h_b[i * arraySize], arraySize * sizeof(float), hipMemcpyHostToDevice));
        }
        
        // Single kernel launch for all sub-tasks
        hipLaunchKernelGGL(masqueradingKernel, dim3(gridSize), dim3(threadsPerBlock), 
                         0, 0, d_a_ptrs, d_b_ptrs, d_c_ptrs, arraySize, numSubTasks);
        CHECK_HIP(hipDeviceSynchronize());
        
        // Copy results back
        for (int i = 0; i < numSubTasks; i++) {
            CHECK_HIP(hipMemcpy(&h_c_masquerading[i * arraySize], d_c[i], arraySize * sizeof(float), hipMemcpyDeviceToHost));
        }
        
        auto end = std::chrono::high_resolution_clock::now();
        double duration = std::chrono::duration<double, std::milli>(end - start).count();
        
        std::cout << "Masquerading execution time: " << duration << " ms" << std::endl;
    }
    
    bool validateResults() {
        std::cout << "Validating results..." << std::endl;
        
        // Check sequential results
        float maxErrorSequential = 0.0f;
        for (int i = 0; i < arraySize * numSubTasks; i++) {
            float error = fabs(h_c_sequential[i] - h_reference[i]);
            maxErrorSequential = fmax(maxErrorSequential, error);
        }
        
        // Check masquerading results
        float maxErrorMasquerading = 0.0f;
        for (int i = 0; i < arraySize * numSubTasks; i++) {
            float error = fabs(h_c_masquerading[i] - h_reference[i]);
            maxErrorMasquerading = fmax(maxErrorMasquerading, error);
        }
        
        std::cout << "Max error (sequential): " << maxErrorSequential << std::endl;
        std::cout << "Max error (masquerading): " << maxErrorMasquerading << std::endl;
        
        // Save results for analysis
        saveResults();
        
        return (maxErrorSequential < 1e-6 && maxErrorMasquerading < 1e-6);
    }
    
    void saveResults() {
        // Save binary arrays for analysis
        std::ofstream seq_file("output/hip_sequential_result.npy", std::ios::binary);
        std::ofstream masq_file("output/hip_masquerading_result.npy", std::ios::binary);
        std::ofstream ref_file("output/hip_reference_result.npy", std::ios::binary);
        
        if (seq_file && masq_file && ref_file) {
            seq_file.write(reinterpret_cast<char*>(h_c_sequential), arraySize * numSubTasks * sizeof(float));
            masq_file.write(reinterpret_cast<char*>(h_c_masquerading), arraySize * numSubTasks * sizeof(float));
            ref_file.write(reinterpret_cast<char*>(h_reference), arraySize * numSubTasks * sizeof(float));
        }
        
        seq_file.close();
        masq_file.close();
        ref_file.close();
    }
    
    void printSummary() {
        std::cout << "\n=== HIP Correctness Validation Summary ===" << std::endl;
        std::cout << "Array size per sub-task: " << arraySize << " elements" << std::endl;
        std::cout << "Number of sub-tasks: " << numSubTasks << std::endl;
        std::cout << "Total elements processed: " << arraySize * numSubTasks << std::endl;
        std::cout << "Memory per array: " << (arraySize * sizeof(float) / 1024.0 / 1024.0) << " MB" << std::endl;
        std::cout << "Total memory: " << (arraySize * numSubTasks * sizeof(float) * 3 / 1024.0 / 1024.0) << " MB" << std::endl;
    }
    
    void cleanup() {
        if (h_a) hipHostFree(h_a);
        if (h_b) hipHostFree(h_b);
        if (h_c_sequential) hipHostFree(h_c_sequential);
        if (h_c_masquerading) hipHostFree(h_c_masquerading);
        if (h_reference) hipHostFree(h_reference);
        
        for (int i = 0; i < numSubTasks; i++) {
            if (d_a[i]) hipFree(d_a[i]);
            if (d_b[i]) hipFree(d_b[i]);
            if (d_c[i]) hipFree(d_c[i]);
        }
        
        if (d_a_ptrs) hipFree(d_a_ptrs);
        if (d_b_ptrs) hipFree(d_b_ptrs);
        if (d_c_ptrs) hipFree(d_c_ptrs);
        
        delete[] d_a;
        delete[] d_b;
        delete[] d_c;
    }
};

int main(int argc, char* argv[]) {
    std::cout << "=== HIP Correctness Validation Framework ===" << std::endl;
    std::cout << "Validating DCU 10x performance improvement claim" << std::endl;
    std::cout << "Target: 128 sub-tasks × 1024×1024 matrices" << std::endl;
    std::cout << std::endl;
    
    try {
        CorrectnessValidator validator;
        validator.printSummary();
        
        // Run validation tests
        validator.runSequential();
        validator.runMasquerading();
        
        // Validate correctness
        bool correct = validator.validateResults();
        
        if (correct) {
            std::cout << "✅ VALIDATION PASSED: All results are correct" << std::endl;
            std::cout << "✅ Both sequential and masquerading modes produce identical results" << std::endl;
        } else {
            std::cout << "❌ VALIDATION FAILED: Results differ significantly" << std::endl;
        }
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}